package com.eonmind.ai.core.draw.schedule;

import com.eonmind.ai.core.draw.schedule.handler.MjSyncProgressHandler;
import com.eonmind.ai.mapper.AiDrawMapper;
import com.eonmind.ai.service.impl.AiDrawPlatformService;
import com.eonmind.common.contract.enums.ai.AiDrwaPlatformTypeEnum;
import com.eonmind.common.provider.entity.ai.AiArtwork;
import com.eonmind.common.provider.entity.ai.AiDrawPlatform;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

@Slf4j
@Component
public class DrawTaskSyncSchedule {
    private static final String syncTaskKey = "sync:draw";
    @Resource
    private RedissonClient redissonClient;

    @Resource
    private AiDrawMapper aiDrawMapper;

    @Resource(name = "aiDefaultTaskExecutor")
    private ThreadPoolTaskExecutor taskExecutor;

    @Resource
    private MjSyncProgressHandler mjSyncProgressHandler;

    @Resource
    private AiDrawPlatformService aiDrawPlatformService;

    @Scheduled(fixedDelay = 1000 * 3)
    public void syncDrawTask() {
        RLock lock = redissonClient.getLock(syncTaskKey);
        boolean success = lock.tryLock();
        if (success) {
            // 进行同步
            // 分页查询待同步的任务
            List<AiArtwork> aiArtworks = aiDrawMapper.selectAllSyncData();
            if (ObjectUtils.isEmpty(aiArtworks)) return;
            else {
                log.info("查询到待同步的绘画任务：{}", aiArtworks.size());
            }
            Map<Long, List<AiArtwork>> workMap = aiArtworks.stream()
                    .collect(Collectors.groupingBy(AiArtwork::getPlatformId));
            CountDownLatch countDownLatch = new CountDownLatch(workMap.size());
            workMap.forEach((platformId, aiArtworkList) -> {
                syncProgress(platformId, aiArtworkList, countDownLatch);
            });
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                log.error("绘画进度同步任务失败：{}", new Date(), e);
                throw new RuntimeException(e);
            }
        }
    }

    private void syncProgress(Long platformId, List<AiArtwork> aiArtworkList, CountDownLatch countDownLatch) {

        taskExecutor.execute(() -> {
            try {
                AiDrawPlatform aiDrawPlatform = aiDrawPlatformService.selectById(platformId);

                AiDrwaPlatformTypeEnum aiDrwaPlatformTypeEnum = AiDrwaPlatformTypeEnum.getByCode(aiDrawPlatform.getPlatformType());
                if (ObjectUtils.isEmpty(aiDrwaPlatformTypeEnum)) {
                    log.error("平台类型错误：{}", aiDrawPlatform.getPlatformType());
                    return;
                }
                switch (aiDrwaPlatformTypeEnum) {
                    case MJ:
                        mjSyncProgressHandler.syncProgress(aiDrawPlatform, aiArtworkList);
                }

            } catch (Exception e) {
                log.error("绘画进度同步任务失败：{}", new Date(), e);
            } finally {
                countDownLatch.countDown();
            }
        });
    }


}
